import face_recognition
import cv2
import numpy as np

# 初始化摄像头
cap = cv2.VideoCapture(0)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)

# 使用HOG模型提高速度
model_for_location = "hog"

# 减少上采样次数
number_of_times_to_upsample = 1

# 已知人员的人脸编码和名字
known_face_encodings = []
known_face_names = []

# 假设你有几个已知人员的图片文件，例如 "person1.jpg", "person2.jpg" 等
images_of_people = ['chenhan.jpg', 'star.jpg']

for image_filename in images_of_people:
    image_file = face_recognition.load_image_file(image_filename)
    encoding = face_recognition.face_encodings(image_file)[0]
    known_face_encodings.append(encoding)
    # 从文件名中提取人名
    name = image_filename.split('.')[0]
    known_face_names.append(name)

# 人脸识别容差值，可以根据需要调整
tolerance = 0.4

try:
    while True:
        ret, frame = cap.read()
        if not ret:
            print("无法读取帧")
            break

        # 转换颜色空间
        rgb_frame = np.ascontiguousarray(frame[:, :, ::-1])  # OpenCV 使用 BGR，而 face_recognition 需要 RGB

        # 检测人脸
        face_locations = face_recognition.face_locations(rgb_frame, model=model_for_location, number_of_times_to_upsample=number_of_times_to_upsample)

        # 如果检测到人脸，则提取面部编码
        if face_locations:
            face_encodings = face_recognition.face_encodings(rgb_frame, face_locations)
            for (top, right, bottom, left), face_encoding in zip(face_locations, face_encodings):
                # 匹配人脸
                matches = face_recognition.compare_faces(known_face_encodings, face_encoding, tolerance)

                name = "Unknown"

                # 如果找到匹配的人脸
                if True in matches:
                    first_match_index = matches.index(True)
                    name = known_face_names[first_match_index]

                # 绘制矩形框
                cv2.rectangle(frame, (left, top), (right, bottom), (0, 0, 255), 2)  # 红色矩形框

                # 定义浅蓝色的BGR值
                light_blue = (255, 191, 216)  # 浅蓝色，这里使用了RGB格式的浅蓝色，OpenCV会自动转换为BGR

                # 将文字颜色设置为浅蓝色，并放大文字
                cv2.putText(frame, name, (left + 6, bottom - 6), cv2.FONT_HERSHEY_DUPLEX, 0.75, light_blue, 2)  # 缩放比例调整为0.75，线条粗细调整为2

        # 显示结果帧
        cv2.imshow('Face Recognition', frame)

        # 退出循环
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

finally:
    # 释放摄像头资源
    cap.release()
    cv2.destroyAllWindows()